package com.springboot.frame.web.csrf;

import org.springframework.util.Assert;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import javax.servlet.http.HttpSession;
import java.util.UUID;

public final class SessionCsrfTokenRepository implements CsrfTokenRepository {

  private static final String SESSION_ATTRIBUTE_NAME = SessionCsrfTokenRepository.class.getName()
      .concat(".CSRF_TOKEN");
  private String parameterName = "_csrf";
  private String headerName = "X-CSRF-TOKEN";

  @Override
  public void saveToken(CsrfToken token, HttpServletRequest request, HttpServletResponse response) {
    if (token == null) {
      HttpSession session = request.getSession(false);
      if (session != null) {
        session.removeAttribute(SESSION_ATTRIBUTE_NAME);
      }
    } else {
      HttpSession session = request.getSession();
      session.setAttribute(SESSION_ATTRIBUTE_NAME, token);
    }
  }

  @Override
  public CsrfToken loadToken(HttpServletRequest request) {
    HttpSession session = request.getSession(false);
    if (session == null) {
      return null;
    }
    return (CsrfToken) session.getAttribute(SESSION_ATTRIBUTE_NAME);
  }

  @Override
  public CsrfToken generateToken(HttpServletRequest request) {
    return new CsrfToken(this.headerName, this.parameterName, UUID.randomUUID().toString());
  }

  public void setParameterName(String parameterName) {
    Assert.hasLength(parameterName, "parameterName cannot be null or empty");
    this.parameterName = parameterName;
  }

  public void setHeaderName(String headerName) {
    Assert.hasLength(headerName, "headerName cannot be null or empty");
    this.headerName = headerName;
  }

}
